package com.tangyuan.ai.util;


import cn.hutool.json.JSONUtil;
import com.tangyuan.ai.entity.dto.AiRequestBodyDTO;
import com.tangyuan.ai.entity.dto.AiRequestImageBodyDTO;
import com.tangyuan.ai.entity.vo.AiImageTaskVo;
import com.tangyuan.ai.entity.vo.ErrorMessageVo;
import com.tangyuan.ai.entity.vo.MessageVo;
import com.tangyuan.ai.entity.vo.UsageVo;
import com.tangyuan.common.util.OkHttpUtils;
import com.tangyuan.common.util.StringUtils;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.json.JSONObject;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import java.util.concurrent.TimeUnit;


@Component
@Slf4j
public class AiUtil {


    @Value("${ai.variables.apiKey:}")
    private String apiKey;

    @Value("${ai.variables.secretKey:}")
    private String secretKey;

    @Value("${ai.variables.imageApiKey:}")
    private String imageApiKey;

    @Value("${ai.variables.imageSecretKey:}")
    private String imageSecretKey;

    @Value("${ai.variables.size:}")
    private String size;

    @Value("${ai.variables.n:}")
    private Integer n;

    @Value("${ai.variables.maxTokens:}")
    private Integer maxTokens;

    @Value("${ai.variables.temperature:}")
    private float temperature;

    @Value("${ai.variables.style:}")
    private String style;



    /**
     * 文生文
     */
    public MessageVo dialogueMessage(AiRequestBodyDTO dto){
        dto.setTemperature(temperature);
        dto.setMaxOutputTokens(maxTokens);
        MediaType mediaType = MediaType.parse("application/json");
        String str = JSONUtil.toJsonStr(dto);
        RequestBody body = RequestBody.create(mediaType, str);
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token="+getAccessToken(apiKey,secretKey))
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .build();
        String bodyResponse = "";
        try {
            Response response = OkHttpUtils.getInstance().newCall(request).execute();
            bodyResponse = response.body().string();
        }catch (Exception e){
            log.error("AI文生文异常:{}",e.toString());
        }
        return exceptionHandling(bodyResponse);
    }

    /**
     * 文生图
     */
    public AiImageTaskVo dialogueImages(AiRequestImageBodyDTO dto) {
        dto.setResolution(size);
        dto.setNum(n);
        dto.setStyle(style);
        MediaType mediaType = MediaType.parse("application/json");
        RequestBody body = RequestBody.create(mediaType, JSONUtil.toJsonStr(dto));
        String token = getAccessToken(imageApiKey, imageSecretKey);
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/rpc/2.0/wenxin/v1/basic/textToImage?access_token="+token)
                .method("POST", body)
                .addHeader("Content-Type", "application/json")
                .build();
        String bodyResponse = "";
        try {
            Response response = OkHttpUtils.getInstance().newCall(request).execute();
            bodyResponse = response.body().string();
        }catch (Exception e){
            log.error("AI文生图异常:{}",e.toString());
        }
        AiImageTaskVo taskVo = null;
        if(bodyResponse.contains("taskId")){
            taskVo = com.alibaba.fastjson.JSONObject.parseObject(bodyResponse, AiImageTaskVo.class);
            Long taskId = taskVo.getData().getTaskId();
            cn.hutool.json.JSONObject taskJson = new cn.hutool.json.JSONObject();
            taskJson.put("taskId",taskId);
            while (true){
                try {
                    TimeUnit.SECONDS.sleep(1);
                    RequestBody imageBody = RequestBody.create(mediaType, taskJson.toString());
                    Request taskRequest = new Request.Builder()
                            .url("https://aip.baidubce.com/rpc/2.0/wenxin/v1/basic/getImg?access_token="+getAccessToken(imageApiKey,imageSecretKey))
                            .method("POST", imageBody)
                            .addHeader("Content-Type", "application/json")
                            .build();
                    Response response = OkHttpUtils.getInstance().newCall(taskRequest).execute();
                    bodyResponse = response.body().string();
                    taskVo = com.alibaba.fastjson.JSONObject.parseObject(bodyResponse, AiImageTaskVo.class);
                    if(taskVo.getData().getStatus() == 1){
                        break;
                    }
                }catch (Exception e){
                    log.error("文生图请求任务异常,taskId:{},异常:{}",taskId,e.toString());
                    taskVo = new AiImageTaskVo();
                }
            }
        }else {
            taskVo = new AiImageTaskVo();
            log.error("文生图请求任务异常,bodyResponse:{}",bodyResponse);
        }
        return taskVo;
    }

    private MessageVo exceptionHandling(String message) {
        MessageVo vo = new MessageVo();
        if(StringUtils.isEmpty(message)){
            vo.setResult("Ai异常");
            vo.setUsage(new UsageVo(0,0,0));
        }else if(message.contains("error_code")){
            ErrorMessageVo messageVo = com.alibaba.fastjson.JSONObject.parseObject(message, ErrorMessageVo.class);
            log.error("Ai请求异常,异常:{}",message);
            vo.setResult(messageVo.getErrorMsg());
            vo.setUsage(new UsageVo(0,0,0));
        }else {
            vo = com.alibaba.fastjson.JSONObject.parseObject(message, MessageVo.class);
        }
        return vo;
    }


    public String getAccessToken(String apiKey,String secretKey) {
        MediaType mediaType = MediaType.parse("application/x-www-form-urlencoded");
        RequestBody body = RequestBody.create(mediaType, "grant_type=client_credentials&client_id=" + apiKey
                + "&client_secret=" + secretKey);
        Request request = new Request.Builder()
                .url("https://aip.baidubce.com/oauth/2.0/token")
                .method("POST", body)
                .addHeader("Content-Type", "application/x-www-form-urlencoded")
                .build();
        try {
            Response response = OkHttpUtils.getInstance().newCall(request).execute();
            return new JSONObject(response.body().string()).getString("access_token");
        }catch (Exception e){
            log.error("请求获取文心 access_token异常,异常:{}",e.toString());
        }
        return "";
    }
}
